"""
Complete Workflow Example for PDEJulia.jl

This example demonstrates the full capabilities of the PDEJulia ecosystem,
showing how all components integrate together for solving PDEs.

# Workflow Overview:
1. Grid creation and boundary conditions
2. Initial condition setup using fields
3. PDE definition and configuration
4. Solver selection and configuration
5. Tracking and visualization setup
6. Simulation execution
7. Data analysis and post-processing
8. Performance benchmarking

# Comparison with py-pde:
This example shows equivalent functionality to py-pde but with Julia's
performance advantages and more flexible type system.
"""

using PDEJulia
using Plots
using BenchmarkTools
using JLD2
using Printf

# =====================================
# 1. BASIC DIFFUSION EQUATION EXAMPLE
# =====================================

println("=== Basic Diffusion Equation ===")

# Create 2D Cartesian grid
grid = CartesianGrid([64, 64], [(-5.0, 5.0), (-5.0, 5.0)])
println("Created grid: ", grid)

# Set periodic boundary conditions
set_boundary_conditions!(grid, [PeriodicBC(), PeriodicBC()])

# Create initial condition: Gaussian blob
u0 = ScalarField(x -> exp(-(x[1]^2 + x[2]^2)), grid)
println("Initial condition statistics: ", get_statistics(u0))

# Define diffusion PDE
pde = DiffusionPDE(diffusivity=0.1)
println("Created PDE: ", pde)

# Setup solver with adaptive time stepping
solver = AdaptiveSolver(alg=Tsit5(), rtol=1e-6, abstol=1e-8)

# Setup tracking
plot_tracker = PlotTracker(interval=1.0, filename="diffusion_evolution.gif")
data_tracker = DataTracker(["norm", "integral", "max_value"])

# Setup storage
storage = FileStorage("diffusion_results.jld2")

# Run simulation
println("Running diffusion simulation...")
@time result = solve(pde, u0, (0.0, 10.0); 
                    solver=solver, 
                    tracker=[plot_tracker, data_tracker], 
                    storage=storage)

println("Simulation completed. Final time: ", result.t)

# =====================================
# 2. ADVANCED PATTERN FORMATION EXAMPLE
# =====================================

println("\n=== Pattern Formation: Swift-Hohenberg Equation ===")

# Create larger grid for pattern formation
grid_2d = CartesianGrid([128, 128], [(0.0, 20.0), (0.0, 20.0)])
set_boundary_conditions!(grid_2d, [PeriodicBC(), PeriodicBC()])

# Initial condition: small random perturbations
u0_sh = ScalarField(x -> 0.1 * (rand() - 0.5), grid_2d)

# Swift-Hohenberg PDE: ∂u/∂t = r*u - (1 + ∇²)²u + u³ - u⁵
pde_sh = SwiftHohenbergPDE(r=0.3, γ=1.0)

# Use implicit solver for stiff equation
solver_sh = ImplicitEulerSolver(dt=0.1, max_steps=1000)

# Track pattern formation
pattern_tracker = PlotTracker(interval=5.0, 
                             filename="pattern_formation.gif",
                             show_colorbar=true)

steady_state_tracker = SteadyStateTracker(rtol=1e-6, interval=10.0)

println("Running pattern formation simulation...")
@time result_sh = solve(pde_sh, u0_sh, (0.0, 200.0);
                       solver=solver_sh,
                       tracker=[pattern_tracker, steady_state_tracker])

if steady_state_tracker.converged
    println("Reached steady state at time: ", steady_state_tracker.convergence_time)
else
    println("Did not reach steady state within time limit")
end

# =====================================
# 3. MULTI-FIELD SYSTEM EXAMPLE
# =====================================

println("\n=== Multi-Field System: Cahn-Hilliard Equation ===")

# Cahn-Hilliard equation for phase separation
# ∂c/∂t = ∇²μ, where μ = -γ∇²c + f'(c)

grid_ch = CartesianGrid([64, 64], [(0.0, 10.0), (0.0, 10.0)])
set_boundary_conditions!(grid_ch, [NeumannBC(), NeumannBC()])

# Initial condition: random composition field
c0 = ScalarField(x -> 0.5 + 0.1 * (rand() - 0.5), grid_ch)

# Create field collection for multi-component system
fields = FieldCollection([c0])  # Would have multiple fields in full system

# Cahn-Hilliard PDE
pde_ch = CahnHilliardPDE(mobility=1.0, interface_width=0.5, gamma=1.0)

# Use Crank-Nicolson for better stability
solver_ch = CrankNicolsonSolver(dt=0.01)

# Track phase separation dynamics
phase_tracker = PlotTracker(interval=2.0, filename="phase_separation.gif")
interface_tracker = CallbackTracker(field -> track_interface_length(field))

println("Running phase separation simulation...")
@time result_ch = solve(pde_ch, c0, (0.0, 50.0);
                       solver=solver_ch,
                       tracker=[phase_tracker, interface_tracker])

# =====================================
# 4. SPHERICAL GEOMETRY EXAMPLE
# =====================================

println("\n=== Spherical Geometry: Reaction-Diffusion ===")

# Spherical grid for radially symmetric problems
grid_sph = SphericalSymGrid(64, (0.1, 10.0))
set_boundary_conditions!(grid_sph, [NeumannBC()])

# Initial condition: localized activation
u0_sph = ScalarField(r -> exp(-((r[1] - 2.0)/0.5)^2), grid_sph)

# Reaction-diffusion with Turing instability
pde_rd = AllenCahnPDE(epsilon=0.1, alpha=1.0, beta=-1.0)

# Explicit solver for fast dynamics
solver_rd = ExplicitEulerSolver(dt=0.001)

# Track radial profiles
profile_tracker = CallbackTracker(field -> extract_radial_profile(field))

println("Running spherical reaction-diffusion...")
@time result_sph = solve(pde_rd, u0_sph, (0.0, 20.0);
                        solver=solver_rd,
                        tracker=profile_tracker)

# =====================================
# 5. PERFORMANCE BENCHMARKING
# =====================================

println("\n=== Performance Benchmarking ===")

# Benchmark different grid sizes
function benchmark_diffusion(n)
    grid = UnitGrid([n, n])
    u0 = ScalarField(x -> exp(-(x[1]^2 + x[2]^2)), grid)
    pde = DiffusionPDE(diffusivity=0.1)
    solver = ExplicitEulerSolver(dt=0.001)
    
    # Benchmark single time step
    @benchmark solve($pde, $u0, 0.1; solver=$solver, tracker=nothing)
end

# Benchmark various grid sizes
grid_sizes = [32, 64, 128, 256]
benchmarks = []

println("Benchmarking different grid sizes:")
for n in grid_sizes
    println("  Grid size: $(n)x$(n)")
    b = benchmark_diffusion(n)
    push!(benchmarks, (n, b))
    
    median_time = median(b.times) * 1e-6  # Convert to ms
    println("    Median time: $(median_time:.2f) ms")
    
    # Memory usage
    memory_mb = median(b.memory) / 1024^2
    println("    Memory usage: $(memory_mb:.2f) MB")
end

# =====================================
# 6. PARAMETER SWEEP EXAMPLE
# =====================================

println("\n=== Parameter Sweep ===")

# Define PDE template
pde_template = params -> DiffusionPDE(diffusivity=params.D)

# Parameter ranges
parameters = (D = [0.01, 0.05, 0.1, 0.5, 1.0],)

# Run parameter sweep
grid_sweep = UnitGrid([32, 32])
u0_sweep = ScalarField(x -> exp(-10*(x[1]-0.5)^2), grid_sweep)

println("Running parameter sweep over diffusivity values...")
@time results_sweep = parameter_sweep(pde_template, parameters, u0_sweep, 5.0;
                                     solver=ExplicitEulerSolver(dt=0.01))

# Analyze sweep results
println("Parameter sweep results:")
for (params, result) in results_sweep
    final_max = maximum(data(result))
    println("  D = $(params.D): final maximum = $(final_max:.4f)")
end

# =====================================
# 7. CUSTOM PDE EXAMPLE
# =====================================

println("\n=== Custom PDE: Kuramoto-Sivashinsky ===")

# Custom PDE implementation
struct CustomKSPDE <: AbstractPDE
    ν::Float64  # Viscosity parameter
    
    CustomKSPDE(; ν=1.0) = new(ν)
end

# Define evolution rate for custom PDE
function evolution_rate(pde::CustomKSPDE, field::ScalarField, t::Float64)
    # KS equation: ∂u/∂t = -∇⁴u - ν∇²u - (∇u)²/2
    u = field
    ∇²u = laplacian(u)
    ∇⁴u = laplacian(∇²u)
    ∇u = gradient(u)
    
    # Nonlinear term: -(∇u)²/2
    nonlinear = -0.5 * dot(∇u, ∇u)  # This would need proper implementation
    
    return -∇⁴u - pde.ν * ∇²u + nonlinear
end

# Test custom PDE
grid_ks = UnitGrid([128])
set_boundary_conditions!(grid_ks, [PeriodicBC()])

u0_ks = ScalarField(x -> sin(2π * x[1]) + 0.1 * sin(4π * x[1]), grid_ks)
pde_ks = CustomKSPDE(ν=1.0)

solver_ks = AdaptiveSolver(rtol=1e-6)
tracker_ks = PlotTracker(interval=1.0)

println("Running custom Kuramoto-Sivashinsky simulation...")
@time result_ks = solve(pde_ks, u0_ks, (0.0, 50.0);
                       solver=solver_ks, tracker=tracker_ks)

# =====================================
# 8. DATA ANALYSIS AND VISUALIZATION
# =====================================

println("\n=== Data Analysis and Visualization ===")

# Load and analyze stored results
stored_data = load("diffusion_results.jld2")
println("Loaded simulation data with $(length(stored_data)) time points")

# Create summary plots
p1 = plot(title="Diffusion Evolution")
for i in 1:5:length(stored_data)
    field_data = stored_data[i]
    contour!(p1, field_data, alpha=0.7)
end

# Time series analysis
tracker_data = get_tracker_data(data_tracker)
p2 = plot(tracker_data.times, tracker_data.norm, 
         title="Field Norm vs Time", xlabel="Time", ylabel="L2 Norm")

# Energy spectrum analysis (for periodic domains)
if has_periodic_bc(grid)
    final_field = stored_data[end]
    spectrum = compute_energy_spectrum(final_field)
    p3 = plot(spectrum.wavenumbers, spectrum.energy,
             title="Energy Spectrum", xlabel="Wavenumber", ylabel="Energy",
             yscale=:log10)
end

# Save analysis results
savefig(p1, "diffusion_evolution.png")
savefig(p2, "norm_evolution.png")

# =====================================
# 9. MIGRATION FROM PY-PDE EXAMPLE
# =====================================

println("\n=== Migration from py-pde ===")

# Show equivalent py-pde code in comments and Julia translation

# Python (py-pde):
# import pde
# grid = pde.UnitGrid([64, 64])
# field = pde.ScalarField.random_uniform(grid, 0.2, 0.3)
# eq = pde.DiffusionPDE(diffusivity=0.1)
# result = eq.solve(field, t_range=10, dt=0.01)

# Julia (PDEJulia):
grid_migration = UnitGrid([64, 64])
field_migration = ScalarField(x -> 0.2 + 0.1 * rand(), grid_migration)
eq_migration = DiffusionPDE(diffusivity=0.1)
result_migration = solve(eq_migration, field_migration, 10.0;
                        solver=ExplicitEulerSolver(dt=0.01))

println("Migration example completed - same functionality, better performance!")

# =====================================
# 10. SUMMARY AND PERFORMANCE REPORT
# =====================================

println("\n=== PERFORMANCE SUMMARY ===")
println("All simulations completed successfully!")
println()
println("Key advantages of PDEJulia.jl:")
println("  • Type-stable, high-performance numerics")
println("  • Flexible grid and field system")
println("  • Integrated visualization and tracking")
println("  • Comprehensive solver suite")
println("  • Easy parameter sweeps and analysis")
println("  • Direct migration path from py-pde")
println()
println("Benchmark summary:")
for (n, b) in benchmarks
    median_time = median(b.times) * 1e-6
    println("  $(n)x$(n) grid: $(median_time:.2f) ms per time step")
end

println("\nWorkflow demonstration complete!")

# Helper functions used in examples
function track_interface_length(field::ScalarField)
    # Simple interface tracking - compute gradient magnitude
    grad = gradient(field)
    interface_strength = norm(grad)
    return Dict("interface_length" => interface_strength)
end

function extract_radial_profile(field::ScalarField)
    # Extract radial profile for spherical grid
    g = grid(field)
    if isa(g, SphericalSymGrid)
        return Dict("radial_profile" => data(field))
    else
        return Dict()
    end
end

function compute_energy_spectrum(field::ScalarField)
    # Compute energy spectrum using FFT
    if has_periodic_bc(grid(field))
        fft_data = fft(data(field))
        energy = abs2.(fft_data)
        wavenumbers = fftfreq(size(field, 1))
        return (wavenumbers=wavenumbers, energy=energy)
    else
        error("Energy spectrum requires periodic boundary conditions")
    end
end

function has_periodic_bc(grid::AbstractGrid)
    bcs = get_boundary_conditions(grid)
    return all(bc -> isa(bc, PeriodicBC), bcs)
end